Super resolution on medical images:¶
Set-up:¶
Imports:¶
In [12]:
%matplotlib inline
%load_ext autoreload
%autoreload 2
import pandas as pd
import fastai
from fastai.vision import *
from fastai.callbacks import *
from fastai.utils.mem import *
import tensorflow as tf
from torchvision.models import vgg16_bn
import torchvision.transforms as transforms
from skimage import measure
from modified_model_fx import *
In [13]:
import warnings
warnings.filterwarnings('ignore')
In [14]:
trans1 = transforms.ToTensor()
trans = transforms.ToPILImage()
Set CUDA:
In [15]:
torch.cuda.set_device(3)
Import data:¶
Input images:
- HR images: (3, 502, 672) original (3, 1004, 1344) cut in 4 pieces
- LR images: (3, 250, 334) original (3, 500, 669) cut in 4 pieces
Paths to data:
In [16]:
path = Path('../../../../../SCRATCH2/marvande/data/train/HR/')
# path to HR data:
path_hr = path / 'HR_patches_train/jpg_images'
# path where we create the LR data:
path_lr = path / 'small-250/train'
# path to MR data (of same size as HR):
path_mr = path / 'small-502/train'
path_mr_test = path / 'small-250/test'
# path to original LR data:
path_lr_or = path/'HR_patches_resized/jpg_images'
assert path.exists(), f"need dataset @ {path}"
assert path_hr.exists()
Have a look at what type of data is in those folders. First, we look at our HR images from path_hr.
In [17]:
il = ImageList.from_folder(path_hr)
ImageList.from_folder(path_hr)
Out[17]:
Create LR and MR data:¶
From this HR data, we create LR versions in path_lr and path_mr for training. LR images will be of size (3, 250, 334) while MR images of the same size as the HR but of lower quality (MR images are used in the second phase of training), thus (3, 502, 672).
In [18]:
def resize_one(fn, i, path, size):
"""resize_one: resizes images to input size and saves them in path,
quality is lowered as to get LR images.
"""
dest = path / fn.relative_to(path_hr)
dest.parent.mkdir(parents=True, exist_ok=True)
img = PIL.Image.open(fn)
img = img.resize(size, resample=PIL.Image.BILINEAR).convert('RGB')
img.save(dest, quality=60)
In [19]:
# create smaller image sets the first time this nb is run:
sets = [(path_lr, (334, 250)), (path_mr, (672, 502))]
for p, size in sets:
if not p.exists():
print(f"resizing to {size} into {p}")
parallel(partial(resize_one, path=p, size=size), il.items)
Have a look at the new LR and MR images and their sizes.
In [20]:
print('LR:')
print(ImageList.from_folder(path_lr))
print('MR:')
print(ImageList.from_folder(path_mr))
Training and validation data:¶
As the model's first phase is on images of the LR size, we create the first batch of training data of size (3, 250, 334). So HR images are downsized to (3, 250, 334) and both HR and LR are transformed using several transformations (see modified_model_fx.py). Those data augmentation techniques help avoid overfitting during training.
In [25]:
# set image size and batch size to which data is transformed:
bs, size = 15, (250, 334)
arch = models.resnet34
src = ImageImageList.from_folder(path_lr).split_by_rand_pct(0.1, seed=42)
src
Out[25]:
Change statistics.
In [22]:
def get_data(bs, size):
"""
get_data: creates training and validation data from LR and HR.
downsizes HR to LR size and applies transformations to both.
"""
#label_from_func: apply func to every input to get its label.
# defining a custom function to extract the labels
data = src.label_from_func(lambda x: path_hr / x.relative_to(path_lr))
#apply data transformations,
data = data.transform(tfms, size=size, tfm_y=True).databunch(bs=bs).normalize(imagenet_stats,do_y=True)
data.c = 3
return data
In [24]:
data = get_data(bs, size)
data
Have a look at a few data from the validation data. Left are the LR while right the corresponding Hr images.
In [ ]:
# lr
x = data.one_batch(ds_type=DatasetType.Valid)[0][4]
# hr
y = data.one_batch(ds_type=DatasetType.Valid)[1][4]
print('LR data shape {} and HR shape {}'.format(list(x.shape), list(y.shape)))
MSE = mse_mult_chann(x.numpy(), y.numpy())
NMSE = measure.compare_nrmse(x.numpy(), y.numpy())
SSIM = ssim_mult_chann(x.numpy(), y.numpy())
print('MSE: {}, NMSE: {}, SSIM: {}'.format(MSE, NMSE, SSIM))
plot_single_image(x, '', (10,10))
plot_single_image(y, '', (10,10))
In [26]:
! cp ../../../../../SCRATCH2/marvande/data/train/HR/small-502/train/models/2a_mod_norm.pth ../../../../../SCRATCH2/marvande/data/train/HR/models/
In [13]:
data.show_batch(ds_type=DatasetType.Valid, rows=2, figsize=(20, 20))
Feature loss:¶
Create loss metrics used in the model.
In [62]:
def gram_matrix(x):
"""
Gram matrix of a set of vectors in an inner
product space is the Hermitian matrix of inner products
"""
n,c,h,w = x.size()
x = x.view(n, c, -1)
return (x @ x.transpose(1,2))/(c*h*w)
Select the data of an image and use it to compute its Gram matrix.
In [63]:
im = data.valid_ds[0][1]
t = im.data
t = torch.stack([t,t])
In [64]:
gram_matrix(t)
Out[64]:
We define a base loss as the L1 loss (F.l1_loss).
In [65]:
base_loss = F.l1_loss
Construct a pre-trained vgg16 model (Very Deep Convolutional Networks for Large-Scale Image Recognition). VGG is another Convolutional Neural Network (CNN) architecture devised in 2014, the 16 layer version is utilised in the loss function for training this model. VGG model. a network pretrained on ImageNet, is used to evaluate the generator model’s loss.
Further, we set requires_grad to False as this is useful when you want to freeze part of your model, or you know in advance that you are not going to use gradients w.r.t. some parameters.
The head of the VGG model is the final layers shown as fully connected and softmax in the above diagram. This head is ignored and the loss function uses the intermediate activations in the backbone of the network, which represent the feature detections.
Those activations can be found by looking through the VGG model to find all the max pooling layers. These are where the grid size changes and features are detected. So we need to select those layers.
In [66]:
vgg_m = vgg16_bn(True).features.cuda().eval()
requires_grad(vgg_m, False)
Select the layer IDs of MaxPool2d blocks:
In [67]:
blocks = [i-1 for i,o in enumerate(children(vgg_m)) if isinstance(o,nn.MaxPool2d)]
blocks, [vgg_m[i] for i in blocks]
Out[67]:
Create the feature loss from the model and layer ids selected above.
Source:
Main points:
- Loss functions using these techniques can be used during the training of U-Net based model architectures
- Feature loss: loss function used is similar to the loss function in the the paper, using VGG-16 but also combined with pixel mean squared error loss loss and gram matrix style loss
- The training of a model can use this loss function based on the VGG model’s activations. The loss function remains fixed throughout the training unlike the critic part of a GAN
- Feature loss: Feature map has 256 channels by 28 by 28. The activations at the same layer for the (target) original image and the generated image are compared using mean squared error or the least absolute error (L1) error for the base loss. These are feature losses. This error function uses L1 error. This allows the loss function to know what features are in the target ground truth image and to evaluate how well the model’s prediction’s features match these rather than only comparing pixel difference. This allows the model being trained with this loss function to produce much finer detail in the generated/predicted features and output.
- Gram matrix style loss: A gram matrix defines a style with respect to specific content. By calculating the gram matrix for each feature activation in the target/ground truth image, it allows the style of that feature to be defined. If the same gram matrix is calculated from the activations of the predictions, the two can be compared to calculate how close the style of the feature prediction is to the target/ground truth image. A gram matrix is the matrix multiplication of the each of the activations and the activation matrix’s transpose. This enables the model to learn and generate predictions of images whose features look correct in their style and in context, with the end result looking more convincing and appear closer or the same as the target/ground truth.
Predictions from models trained with this loss function: The generated predictions from trained models using loss functions based on these techniques have both convincing fine detail and style. That style and fine detail may be different aspects of image quality be predicting fine pixel detail or the predicting correct colours.
In [68]:
class FeatureLoss(nn.Module):
def __init__(self, m_feat, layer_ids, layer_wgts):
super().__init__()
self.m_feat = m_feat
self.loss_features = [self.m_feat[i] for i in layer_ids]
self.hooks = hook_outputs(self.loss_features, detach=False)
self.wgts = layer_wgts
self.metric_names = ['pixel',] + [f'feat_{i}' for i in range(len(layer_ids))
] + [f'gram_{i}' for i in range(len(layer_ids))]
def make_features(self, x, clone=False):
self.m_feat(x)
return [(o.clone() if clone else o) for o in self.hooks.stored]
def forward(self, input, target):
out_feat = self.make_features(target, clone=True)
in_feat = self.make_features(input)
self.feat_losses = [base_loss(input,target)]
#feat losses
self.feat_losses += [base_loss(f_in, f_out)*w
for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
#gram:
self.feat_losses += [base_loss(gram_matrix(f_in), gram_matrix(f_out))*w**2 * 5e3
for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
self.metrics = dict(zip(self.metric_names, self.feat_losses))
return sum(self.feat_losses)
def __del__(self): self.hooks.remove()
In [69]:
feat_loss = FeatureLoss(vgg_m, blocks[2:5], [5,15,2])
Train¶
In [73]:
lr = 1e-3
def do_fit(save_name, lrs=slice(lr), pct_start=0.9):
learn.fit_one_cycle(10, lrs, pct_start=pct_start)
learn.save(save_name)
learn.show_results(rows=1, imgsize=10)
In [72]:
# delete all tensors and free cache:
for obj in gc.get_objects():
if torch.is_tensor(obj):
del obj
torch.cuda.empty_cache()
gc.collect()
#learn.destroy()
#get free memory (in MBs) for the currently selected gpu id, after emptying the cache
print(
'free memory (in MBs) for the currently selected gpu id, after emptying the cache: ',
gpu_mem_get_free_no_cache())
print(
'used memory (in MBs) for the currently selected gpu id, after emptying the cache:',
gpu_mem_get_used_no_cache())
gpu_mem_get_all()
Out[72]:
In [74]:
wd = 1e-3
learn = unet_learner(data,
arch,
wd=wd,
loss_func=feat_loss,
callback_fns=LossMetrics,
blur=True,
norm_type=NormType.Weight)
# garbage collection:
gc.collect()
Out[74]:
In [75]:
learn.lr_find()
learn.recorder.plot()
In [55]:
do_fit('1a_mod_norm', slice(lr*10))
In [57]:
learn.unfreeze()
In [58]:
do_fit('1b_mod_norm', slice(1e-5,lr))
Phase 2:
In [27]:
torch.cuda.empty_cache()
gc.collect()
#learn.destroy()
#get free memory (in MBs) for the currently selected gpu id, after emptying the cache
print(
'free memory (in MBs) for the currently selected gpu id, after emptying the cache: ',
gpu_mem_get_free_no_cache())
print(
'used memory (in MBs) for the currently selected gpu id, after emptying the cache:',
gpu_mem_get_used_no_cache())
gpu_mem_get_all()
Out[27]:
In [28]:
new_size = (502, 672)
bs = 4
data = get_data(bs,new_size)
wd = 1e-3
learn = unet_learner(data, arch, wd=wd, loss_func=feat_loss, callback_fns=LossMetrics,
blur=True, norm_type=NormType.Weight)
lr = 1e-3
# garbage collection:
gc.collect();
In [29]:
learn.data = data
learn.freeze()
gc.collect()
Out[29]:
In [30]:
data
Out[30]:
In [31]:
learn.load('1b_mod_norm');
In [32]:
do_fit('2a_mod_norm')
In [33]:
learn.unfreeze()
In [34]:
do_fit('2b_mod_norm', slice(1e-6,1e-4), pct_start=0.3)
Testing :¶
Reconstruct the learn object, this time with l1_loss and not with feature_loss, why?
In [23]:
learn = None
gc.collect()
learn = unet_learner(data,
arch,
loss_func=F.l1_loss,
blur=True,
norm_type=NormType.Weight)
Prediction:¶
Set the sizes of the LR input images (size_lr) for testing and the HR images (size_mr).
In [24]:
size_mr = (3, 502, 672)
size_lr = (3, 250, 334)
path_test = path/'HR_patches_test/cut_images/jpeg_images/'
path_mr_test = path / 'small-250/test'
#Check free GPU RAM:
free = gpu_mem_get_free_no_cache()
print(f"using size={size}, have {free}MB of GPU RAM free")
In [25]:
# Create test images:
test_files = []
with os.scandir(path_test) as entries:
test_files = [entry.name for entry in entries if entry.is_file()]
size = (334, 250)
for file in test_files:
if not path_test.exists():
img = PIL.Image.open(path_test/file)
img = img.resize(size, resample=PIL.Image.BILINEAR).convert('RGB')
img.save(path_mr_test/file, quality=60)
Create testing data, data_mr of size HR and data_lr of size LR:
In [26]:
data_mr = (ImageImageList.from_folder(path_test).split_by_rand_pct(0.1, seed=42)
.label_from_func(lambda x: path_test/x.name)
.transform(tfms, size=size_mr, tfm_y=True)
.databunch(bs=1).normalize(imagenet_stats, do_y=True))
data_mr.c = 3
data_lr = (ImageImageList.from_folder(path_mr_test).split_by_rand_pct(0.1, seed=42)
.label_from_func(lambda x: path_test/x.name)
.transform(tfms, size=size_lr, tfm_y=True)
.databunch(bs=1).normalize(imagenet_stats, do_y=True))
data_lr.c = 3
Load the learn object from the last phase (2b_mod) and set it's data to data_mr (why?).
In [27]:
learn.load('2b_mod_norm')
learn.data = data_mr
Select the images we are going to use for testing from data_mr and data_lr:
In [28]:
# Ground truth HR image:
gt_HR = data_mr.valid_ds.y.items[1]
# LR version of the same image:
lr = data_lr.valid_ds.x.items[1]
To be able to apply the model and predict the HR image from the LR, we need to resample the LR image to be of the same size as the HR, as the model inputs and outputs images of the same size.
In [29]:
lr_data = open_image(lr)
# resample to same size as HR:
# for pytorch, have to add a first new dimension to indicate bs = 1
# now lr_data of shape [1, 3, 250, 334]
lr_data = lr_data.data.unsqueeze(0)
lr_resized = torch.nn.functional.interpolate(lr_data, size_mr[1:],
mode='bicubic')
# remove the previous added dimension
lr_resized = lr_resized.squeeze()
# plot resized LR:
plot_single_image(lr_resized,
'Resized LR from size {} to size {}'\
.format(list(lr_data.squeeze().shape),
list(lr_resized.shape)), (10,10))
Now that we have resized the LR to the same size as the ground truth HR, we can feed it to the model and predict a new HR image.
In [30]:
# Ground truth HR:
im_HR_gt = open_image(gt_HR)
#print('HR ground thruth shape: {}'.format(list(im_HR_gt.shape)))
# Prediction of model:
p, img_pred, b = learn.predict(Image(lr_resized))
#print('Reconstructed HR shape: {}'.format(list(p.shape)))
# Assert reconstructed HR has same shape as ground truth HR:
assert(list(p.shape) == list(im_HR_gt.shape))
Visualisation:¶
Compare this predicted image to the ground truth HR and LR:
In [31]:
plot_3_images(lr_resized, trans1(trans(im_HR_gt.data)).numpy(), img_pred, (20,20))
In [43]:
print('PRed')
print(np.max(img_pred.numpy()))
print(np.min(img_pred.numpy()))
print('Ground truth HR')
print(np.max(trans1(trans(im_HR_gt.data)).numpy()))
print(np.min(trans1(trans(im_HR_gt.data)).numpy()))
print('LR bic')
print(np.max(lr_resized.numpy()))
print(np.min(lr_resized.numpy()))
min(max_HR, trans1(trans(im_HR_gt.data)).numpy())
max(min_HR, trans1(trans(im_HR_gt.data)).numpy())
In [47]:
max_HR = np.max(trans1(trans(im_HR_gt.data)).numpy())
min_HR = np.min(trans1(trans(im_HR_gt.data)).numpy())
test_im = np.clip(a_max = max_HR,a_min = min_HR, a = img_pred.numpy())
print(mse_mult_chann(img_pred.numpy(), trans1(trans(im_HR_gt.data)).numpy()))
mse_mult_chann(test_im, trans1(trans(im_HR_gt.data)).numpy())
Out[47]:
Evaluate loss and error metrics:¶
In [39]:
compare_images_metrics(img_pred.numpy(),trans1(trans(im_HR_gt.data)).numpy(),lr_resized.numpy(),'')
In [34]:
p = '../../../../../SCRATCH2/marvande/data/train/HR/small-250/test/0131_[46833,16191]_part_1_-3.jpg'
pattern = "test\/(.*?)_"
substring = re.search(pattern, p).group(1)
substring
Out[34]:
In [35]:
## Plot for several:
LR_list = ImageList.from_folder(path_mr_test).items
HR_list = ImageList.from_folder(path_test).items
phenotypes = {'0': 'CD4', '1': 'CK','2': 'DAPI', '3': 'CD3', '4': 'FoxP3', '5': 'CD8'}
for i in range(len(LR_list[0:5])):
lr_data = open_image(LR_list[i])
# get file string
pattern = "test\/(.*?)\.jpg"
substring = re.search(pattern, str(LR_list[i])).group(1)
file_name = substring+'.jpg'
# get patient number:
pattern = "test\/(.*?)_"
patient = re.search(pattern, str(LR_list[i])).group(1)
# get phenotype number:
pattern2 = "_\-(.*?)\.jpg"
phenotype_numb = re.search(pattern2,file_name).group(1)
phenotype = phenotypes[phenotype_numb]
# get location number:
pattern2 = "\[(.*?)\]"
location = re.search(pattern2,file_name).group(1)
# resample to same size as HR:
# for pytorch, have to add a first new dimension to indicate bs = 1
# now lr_data of shape [1, 3, 250, 334]
lr_data = lr_data.data.unsqueeze(0)
lr_resized = torch.nn.functional.interpolate(lr_data, size_mr[1:],
mode='bicubic')
# remove the previous added dimension
lr_resized = lr_resized.squeeze()
# corresponding ground truth HR:
im_HR_gt = open_image(HR_list[i])
#print('HR ground thruth shape: {}'.format(list(im_HR_gt.shape)))
# Prediction of model:
p, img_pred, b = learn.predict(Image(lr_resized))
#print('Reconstructed HR shape: {}'.format(list(p.shape)))
# Assert reconstructed HR has same shape as ground truth HR:
assert(list(p.shape) == list(im_HR_gt.shape))
gt_HR_nd = trans1(trans(im_HR_gt.data)).numpy()
pred_HR_nd = img_pred.numpy()
compare_images_metrics(pred_HR_nd,gt_HR_nd, 'phenotype: {}, patient: {}, location: [{}]'.format(phenotype, patient, location))
